import os
import numpy as np
import torchvision.transforms as T
from PIL import Image
import onnxruntime as ort
from omnidet.main import collect_tupperware

def inference(args):
    input_folder = '/home/li/1_slimai_debug/find_dataset_py/detection_images/input'  # 输入图片文件夹
    output_folder = '/home/li/1_slimai_debug/find_dataset_py/detection_images/output'  # 输出图片文件夹
    # 初始化 ONNX Runtime
    ort_session = ort.InferenceSession('/home/li/深度学习/classid_train/git_5.22/omnidet/modified_model.onnx')

    # 遍历输入文件夹中的所有图片
    for img_name in os.listdir(input_folder):
        img_path = os.path.join(input_folder, img_name)
        # 加载图片
        image = Image.open(img_path).convert("RGB")
        originalImage = image
        cropped_coords = dict(Car1=dict(FV=(114, 110, 1176, 610),
                                        MVL=(343, 5, 1088, 411),
                                        MVR=(185, 5, 915, 425),
                                        RV=(186, 203, 1105, 630)),
                              Car2=dict(FV=(160, 272, 1030, 677),
                                        MVL=(327, 7, 1096, 410),
                                        MVR=(175, 4, 935, 404),
                                        RV=(285, 187, 1000, 572)),
                              Car3=dict(FV=(300, 20, 1620, 800),
                                        TEST=(200, 20, 1000, 600)))
        cropped_coords = cropped_coords["Car1"]["FV"]
        image = image.crop(cropped_coords)
        cutimage = image
        to_tensor = T.ToTensor()
        resize = T.Resize((args.input_height, args.input_width),
                           interpolation=T.InterpolationMode.BICUBIC)
        image = resize(image)
        image_tensor = to_tensor(image)
        input_image = image_tensor.unsqueeze(0)

        # 转换为 NumPy 数组并准备 ONNX 模型的输入
        input_image_np = input_image.numpy()

        # 使用 ONNX Runtime 进行推理
        ort_inputs = {ort_session.get_inputs()[0].name: input_image_np}
        ort_outs = ort_session.run(None, ort_inputs)

        # 获取目标检测结果
        predictions = ort_outs[0]

        # 获取分类的最大值索引
        predicted_class = np.argmax(predictions, axis=1)

        print(f"Predicted class: {predicted_class}")
        print(f"Prediction probabilities: {predictions}")
        return predicted_class[0], predictions




if __name__ == '__main__':
    args = collect_tupperware()
    inference(args)